﻿#include "PullStream.h"

#include <QDebug>
#include <QFile>
#include <QCoreApplication>
#include <QJsonObject>
#include <QJsonDocument>
#include <QJsonArray>
#include <QImage>
#include <QtEndian>

PullStream::PullStream(QObject *parent)
    : QThread(parent)
    , bInit(false)
{
    initConfig();

    start();

}

void PullStream::initConfig()
{
    QString filePath = QCoreApplication::applicationDirPath() + "/config.json";
    QFile file(filePath);
    if(file.open(QIODevice::ReadOnly))
    {
        QJsonObject json = QJsonDocument::fromJson(file.readAll()).object();
        rtmpAddr = json["addr"].toString();
    }
}

void PullStream::freeMemory()
{
    if(buffer)
    {
        av_free(buffer);
    }
    if(pFrame)
    {
        av_frame_free(&pFrame);
    }
    if(pFrameRGB)
    {
        av_frame_free(&pFrameRGB);
    }
    if(pCodecCtx)
    {
        avcodec_free_context(&pCodecCtx);
    }
    if(pFormatCtx)
    {
        avformat_free_context(pFormatCtx);
    }
}

void PullStream::parserFlvSEI(const AVPacket * const packet)
{
    uint8_t* ptr = packet->data;
    uint8_t* ptr_end = packet->data + packet->size;
    while(ptr < ptr_end)
    {
        uint32_t size = qFromBigEndian<uint32_t>(ptr);
        ptr += 4;
        uint8_t type = *ptr;
        if((type & 0x1f) == 0x06)
        {
            // SEI数据
            //QByteArray sei_data = QByteArray((char*)ptr, size);
            //qDebug() << "sei: "<< sei_data.toHex(' ');

            // 解析SEI数据
            // 1.1 跳过0X06 05
            ptr += 2;

            // 1.2 解析payload大小
            int payload_size = 0;
            uint8_t byte = 0;
            do {
                byte = *ptr++;
                payload_size += byte;
            } while (byte == 0xFF);

            // payload_size - uuid(16) - 1(0x80)
            int sei_data_size = payload_size - 16 - 1;

            QByteArray uuid = QByteArray((char*)ptr, 16);
            ptr += 16;
            QByteArray sei_data = QByteArray((char*)ptr, sei_data_size);
            ptr += sei_data_size;

//            qDebug() << "uuid: " << uuid.toHex(' ');

//            qDebug() << "sei_data: " << sei_data.toHex(' ');

//            qDebug() << "end flag: "<< QString::number(*ptr, 16);
            emit emitUuidSEI(uuid, sei_data);

            return;
        }else{
            ptr += size;
        }
    }
}

void PullStream::run()
{
    while(true)
    {
        if(bInit)
        {
            if (av_read_frame(pFormatCtx, pPacket) >= 0) {
                if (pPacket->stream_index == videoStream) {

                    // 从关键帧中解析SEI
                    if((pPacket->flags & AV_PKT_FLAG_KEY) != 0)
                    {
                        qDebug() << "idr";
                        //qDebug() << "size: " << pPacket->size << "data: " << QByteArray((char*)pPacket->data,pPacket->size).left(500).toHex(' ');
                        parserFlvSEI(pPacket);
                    }


                    // 解码视频帧
                    int response = avcodec_send_packet(pCodecCtx, pPacket);
                    if (response < 0) {
                        qDebug() << "Error sending a packet for decoding";
                        return;
                    }

                    while (response >= 0) {
                        response = avcodec_receive_frame(pCodecCtx, pFrame);
                        if (response == AVERROR(EAGAIN) || response == AVERROR_EOF) {
                            break;
                        } else if (response < 0) {
                            qDebug() << "Error while decoding";
                            return;
                        }

                        // 转换图像格式为RGB
                        sws_scale(sws_ctx, (uint8_t const * const *)pFrame->data,
                                  pFrame->linesize, 0, pCodecCtx->height,
                                  pFrameRGB->data, pFrameRGB->linesize);

                        // 创建QImage并显示
                        QImage img(pFrameRGB->data[0], pCodecCtx->width, pCodecCtx->height,
                                pFrameRGB->linesize[0], QImage::Format_RGB888);

                        emit emitImage(img);
                    }
                }
                av_packet_unref(pPacket);
            }
        }
        else
        {
            QThread::msleep(5);
        }
    }
}

bool PullStream::init()
{
    bInit = false;
    if(rtmpAddr.isEmpty())
    {
        qDebug() << __FILE__<< __FUNCTION__ << __LINE__ << "rtmp addr is null!";
        return false;
    }

    pFormatCtx = avformat_alloc_context();
    if(!pFormatCtx)
    {
        return false;
    }

    // 打开RTMP流
    if (avformat_open_input(&pFormatCtx, rtmpAddr.toStdString().c_str(), nullptr, nullptr) != 0)
    {
        freeMemory();
        qDebug() << __FILE__<< __FUNCTION__ << __LINE__ << "Could not open input stream.";
        return false;
    }

    // 获取流信息
    if (avformat_find_stream_info(pFormatCtx, nullptr) < 0) {
        freeMemory();
        qDebug() << __FILE__<< __FUNCTION__ << __LINE__ << "Could not find stream information.";
        return false;
    }

    // 查找视频流
    for (unsigned int i = 0; i < pFormatCtx->nb_streams; i++) {
        if (pFormatCtx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            videoStream = i;
            break;
        }
    }

    if (videoStream == -1) {
        qDebug() << __FILE__<< __FUNCTION__ << __LINE__ << "Could not find video stream.";
        return false;
    }

    // 获取编解码器参数和编解码器
    AVCodecParameters *pCodecPar = pFormatCtx->streams[videoStream]->codecpar;
    pCodec = (AVCodec*)avcodec_find_decoder(pCodecPar->codec_id);
    if (pCodec == nullptr) {
        freeMemory();
        qDebug() << __FILE__<< __FUNCTION__ << __LINE__ << "Codec not found.";
        return false;
    }

    pCodecCtx = avcodec_alloc_context3(pCodec);
    if(!pCodecCtx)
    {
        freeMemory();
        qDebug() << __FILE__<< __FUNCTION__ << __LINE__ << "Codec not found.";
        return false;
    }

    if(avcodec_parameters_to_context(pCodecCtx, pCodecPar) < 0)
    {
        freeMemory();
        qDebug() << __FILE__<< __FUNCTION__ << __LINE__ << "avcodec_parameters_to_context error!";
        return false;
    }

    // 打开编解码器
    if (avcodec_open2(pCodecCtx, pCodec, nullptr) < 0)
    {
        freeMemory();
        qDebug() << __FILE__<< __FUNCTION__ << __LINE__ << "avcodec_open2 error.";
        return false;
    }

    // 分配帧
    pFrame = av_frame_alloc();
    pFrameRGB = av_frame_alloc();
    if (pFrameRGB == nullptr || pFrame == nullptr)
    {
        freeMemory();
        qDebug() << __FILE__<< __FUNCTION__ << __LINE__ << "Could not alloc AVFrame.";
        return false;
    }

    // 确定所需的缓冲区大小并分配缓冲区
    int numBytes = av_image_get_buffer_size(AV_PIX_FMT_RGB24, pCodecCtx->width, pCodecCtx->height, 1);
    buffer = (uint8_t *)av_malloc(numBytes * sizeof(uint8_t));
    if(!buffer)
    {
        freeMemory();
        qDebug() << __FILE__<< __FUNCTION__ << __LINE__ << "Could not alloc AVFrame.";
        return false;
    }

    // 将缓冲区与pFrameRGB关联
    av_image_fill_arrays(pFrameRGB->data, pFrameRGB->linesize, buffer,
                         AV_PIX_FMT_RGB24, pCodecCtx->width, pCodecCtx->height, 1);

    // 初始化SWS上下文用于图像转换
    sws_ctx = sws_getContext(pCodecCtx->width, pCodecCtx->height, pCodecCtx->pix_fmt,
                             pCodecCtx->width, pCodecCtx->height, AV_PIX_FMT_RGB24,
                             SWS_BILINEAR, nullptr, nullptr, nullptr);

    if(!sws_ctx)
    {
        freeMemory();
        qDebug() << __FILE__<< __FUNCTION__ << __LINE__ << "Could not alloc sws_ctx.";
    }

    pPacket = av_packet_alloc();
    if(!pPacket)
    {
        freeMemory();
        qDebug() << __FILE__<< __FUNCTION__ << __LINE__ << "Could not alloc AVFrame.";
    }

    bInit = true;

    return true;
}
